#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简单备份工具 - 兼容 Python 3.8+
直接从 SQLite 数据库备份数据
"""

import os
import sys
import json
import shutil
import sqlite3
from pathlib import Path
from datetime import datetime

def backup_now():
    """立即创建备份"""
    print("="*60)
    print("🔄 ChromaDB 数据备份工具")
    print("="*60)
    print()
    
    # 路径设置
    project_root = Path(__file__).parent.parent
    chroma_path = project_root / 'chroma'
    backup_dir = project_root / 'chroma_backups' / 'manual'
    backup_dir.mkdir(parents=True, exist_ok=True)
    
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    
    # 检查数据库文件
    db_file = chroma_path / 'chroma.sqlite3'
    if not db_file.exists():
        print(f"❌ 数据库文件不存在: {db_file}")
        return False
    
    print(f"📁 数据库文件: {db_file}")
    
    # 1. 直接复制数据库文件（最安全）
    print()
    print("💾 备份数据库文件...")
    try:
        backup_folder = backup_dir / f"chroma_{timestamp}"
        shutil.copytree(chroma_path, backup_folder)
        
        # 计算备份大小
        total_size = sum(f.stat().st_size for f in backup_folder.rglob('*') if f.is_file())
        print(f"✅ 数据库文件备份成功: {backup_folder.name}")
        print(f"   大小: {total_size/1024:.1f} KB")
        print(f"   位置: {backup_folder}")
    except Exception as e:
        print(f"❌ 数据库备份失败: {e}")
        return False
    
    # 2. 尝试导出 JSON（如果失败不影响主备份）
    print()
    print("📥 尝试导出数据为 JSON...")
    try:
        conn = sqlite3.connect(str(db_file))
        cursor = conn.cursor()
        
        # 查询记录数
        cursor.execute("SELECT COUNT(*) FROM embeddings;")
        count = cursor.fetchone()[0]
        print(f"   数据库中有 {count} 条记录")
        
        if count > 0:
            # 读取所有记录
            cursor.execute("SELECT * FROM embeddings;")
            rows = cursor.fetchall()
            
            # 读取 metadata
            cursor.execute("SELECT id, key, string_value FROM embedding_metadata;")
            metadata_rows = cursor.fetchall()
            
            # 组织 metadata
            metadata_dict = {}
            for row in metadata_rows:
                record_id, key, value = row
                if record_id not in metadata_dict:
                    metadata_dict[record_id] = {}
                metadata_dict[record_id][key] = value
            
            # 构建导出数据
            export_data = {
                'backup_time': datetime.now().isoformat(),
                'total_records': count,
                'conversations': []
            }
            
            for row in rows:
                record_id = row[2]  # embedding_id
                export_data['conversations'].append({
                    'id': record_id,
                    'metadata': metadata_dict.get(record_id, {}),
                    'note': 'Full data in database backup'
                })
            
            conn.close()
            
            # 写入 JSON
            json_file = backup_dir / f"backup_{timestamp}.json"
            with open(json_file, 'w', encoding='utf-8') as f:
                json.dump(export_data, f, ensure_ascii=False, indent=2)
            
            file_size = json_file.stat().st_size / 1024
            print(f"✅ JSON 备份成功: {json_file.name}")
            print(f"   大小: {file_size:.1f} KB")
        else:
            print("   ⚠️  数据库为空，跳过 JSON 导出")
            conn.close()
            
    except Exception as e:
        print(f"   ⚠️  JSON 导出失败（数据库文件备份已完成）: {e}")
    
    print()
    print("="*60)
    print(f"✅ 备份完成！")
    print(f"📁 备份位置: {backup_dir}")
    print()
    print("💡 提示：")
    print(f"   - 数据库文件备份: {backup_folder.name}")
    print(f"   - 要恢复数据，只需将备份文件夹复制回 chroma 目录")
    print("="*60)
    
    return True

if __name__ == '__main__':
    try:
        backup_now()
    except KeyboardInterrupt:
        print("\n\n👋 已取消")
    except Exception as e:
        print(f"\n❌ 错误: {e}")
        import traceback
        traceback.print_exc()
